from langchain.chat_models import ChatOpenAI
from openai import OpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.callbacks import get_openai_callback



def predict_openai(model,instruction,text):
    """

    :param model:
    :param instruction:
    :param text:
    :return:
    """
    response = model.chat.completions.create(
        model="gpt-3.5-turbo",
        messages = [
            {"role":"system","content":"You are a helpful assistant."},
            {"role":"user","content":instruction + text}
        ]
    )
    result = response.choices[0].message.content
    return result

if __name__ =="__main__":
    instruction = "你现在是一个医疗对话要素抽取专家。\n" \
                  "请针对下面对话内容抽取药品名称、药物类别、医疗检查、医疗操作、现病史、辅助检查、诊断报告和医疗建议等内容，" \
                  "并且以JSON格式返回，Key为上述待抽取的字段名称，value为抽取内容\n" \
                  "注意事项：（1）药品名称、药物类别、医疗检查和医疗操作内容会在多个对话中存在，因此value要以list存在，若抽取为空，则为一个空的list"
    text = '治疗情况： 时间：0000-00-00至0000-00-00 医院科室：未填写未填写 治疗过程：脑垂体瘤手术痊愈 时间：0000-00-00至0000-00-00 医院科室：未填写未填写 治疗过程：脑垂体瘤术后痊愈用药情况：阿司匹林肠溶缓释片，卡维地洛片，瑞舒伐他丁钙片因动脉硬化，现服：阿司匹林肠溶缓释片50mg1次，早饭后，卡维地洛片5mg1次，早饭后，瑞舒伐他丁钙片5mg1次，晚饭后，服用3-4年了。胃药未服。现时常头晕，无力，胀气，胃灼热。'
    client = OpenAI(
        api_key='hk-mtiquv1000041663a49a34520ad3294132cbce7abf1c2ef3',
        base_url="https://api.openai-hk.com/v1"
    )
    result=predict_openai(model=client,instruction=instruction,text=text)
    print(result)